// 实验2：pid替换成阻抗模型，mbk：1/40/0，k设置为0直接到达接触点 接触阶段+td模式
#include <rclcpp/rclcpp.hpp>

// MoveIt Servo
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/planning_scene_monitor/planning_scene_monitor.h>

#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <tf2_ros/buffer.h>
#include <rcl_interfaces/msg/set_parameters_result.hpp>

#include <atomic>
#include <algorithm>
#include <cmath>
#include <string>
#include <chrono>

using namespace std::chrono_literals;

namespace imp_node {  // 避免与其他源文件的全局符号冲突

// ========== 小工具 ==========
static inline double clamp_scalar(double v, double lo, double hi) {
  return v < lo ? lo : (v > hi ? hi : v);
}
static inline double sgn(double x) { return (x > 0) - (x < 0); }

// ========== 节点句柄/通信 ==========
static rclcpp::Node::SharedPtr node_;
static rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_pub_;
static rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;
static rclcpp::node_interfaces::OnSetParametersCallbackHandle::SharedPtr param_cb_handle;

static const rclcpp::Logger LOGGER = rclcpp::get_logger("plane_impedance");

// ========== 参数（可在线调） ==========
// 控制轴向：0=x, 1=y, 2=z
static std::atomic<int>    g_axis{0};

// 目标恒力（N）与线性映射增益 κ：u = κ * (F_target - F_meas)
static std::atomic<double> g_f_target{10.0};
static std::atomic<double> g_kappa{1.0};

// 阻抗参数：m_d, b_d, k_d
static std::atomic<double> g_md{1.0};
static std::atomic<double> g_bd{200.0};
static std::atomic<double> g_kd{0.0}; // 没有接触前直接设为0

// 限幅：速度/加速度
static std::atomic<double> g_max_vel{0.4};   // m/s
static std::atomic<double> g_max_acc{1.5};   // m/s^2

// Twist frame（需与 Servo 配置一致）
static std::string g_twist_frame{"base_link"};

// ========== 状态变量 ==========
static std::atomic<double> g_force_meas{0.0}; // 当前测得力（示例取 Fz）
static double x_ = 0.0;        // 位移修正（m）
static double xdot_ = 0.0;     // 位移修正速度（m/s）
static double u_prev_ = 0.0;   // 上周期控制量（便于调试打印）
static rclcpp::Time t_prev_;
static bool has_prev_time_ = false;

// ======================
// >>> TD/CONTACT：新增 —— 接触判定与 TD 软启动
// ======================

// 接触判定阈值（带滞回 + 力斜率）
static std::atomic<double> g_f_enter{0.5};    // 进入接触阈值 |Fe| >= f_enter 或 |dF/dt| >= df_enter
static std::atomic<double> g_f_exit{0.3};     // 退出接触阈值 |Fe| <= f_exit 且 |dF/dt| 小
static std::atomic<double> g_df_enter{15.0};  // N/s 力斜率阈值（进入接触）

static bool   in_contact_ = false;
static bool   has_prevF_  = false;
static double Fe_prev_    = 0.0;
static rclcpp::Time tF_prev_;

// TD 参数与状态（与 PID 版保持同名/同逻辑）
static std::atomic<bool>   g_td_enable{true};   // 是否启用 TD（仅在接触期启用）
static std::atomic<double> g_td_r0{8.0};        // TD“最速”参数 r0（接触瞬间 5~15 比较温和）
static double g_td_v1 = 0.0;                    // TD 状态 v1（平滑后的设定力）
static double g_td_v2 = 0.0;                    // TD 状态 v2（导数）
static bool   g_td_inited = false;              // TD 是否已初始化

// Han 型 TD 的核心“fhan”（与 PID 代码一致）
static inline double fhan(double x1_minus_ref, double x2, double r, double h)
{
  const double d  = r * h * h;
  const double a0 = h * x2;
  const double y  = x1_minus_ref + a0;
  const double abs_y = std::fabs(y);
  const double a1 = std::sqrt(d * (d + 8.0 * abs_y));
  const double a2 = a0 + sgn(y) * (a1 - d) * 0.5;
  const double a  = a2;
  double f = -r * ( (std::fabs(a) > d) ? sgn(a) : (a / d) );
  return f;
}
static inline void td_update(double Ref, double r0, double h, double &v1, double &v2)
{
  const double f0 = fhan(v1 - Ref, v2, r0, h);
  v1 += h * v2;
  v2 += h * f0;
}

// ========== 订阅 /sim_force ==========
static void forceCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
{
  // 本示例以 Z 向为被控量；若要改 X/Y，请在此切换为 msg->wrench.force.x / .y
  g_force_meas.store(msg->wrench.force.z, std::memory_order_relaxed);

  RCLCPP_INFO_THROTTLE(
      LOGGER, *node_->get_clock(), 500,
      "[/sim_force] Fz=%.3f N (target=%.2f)",
      msg->wrench.force.z, g_f_target.load());
}

// ========== 定时器回调：阻抗模型输出修正量 ==========
static void controlTimer()
{
  // 计算 dt
  const rclcpp::Time now = node_->now();
  double dt = 0.01; // 默认 10ms
  if (has_prev_time_) {
    dt = std::max(1e-4, (now - t_prev_).seconds());
  }
  t_prev_ = now;
  has_prev_time_ = true;

  // 读取力
  const double Fe   = g_force_meas.load(std::memory_order_relaxed);
  const double Fref = g_f_target.load(std::memory_order_relaxed);
  const double kappa = g_kappa.load(std::memory_order_relaxed);

  // ======================
  // >>> TD/CONTACT：接触判定（基于 |Fe| 与 |dF/dt|）
  // ======================
  double dFdt = 0.0;
  if (has_prevF_) {
    double dts = std::max(1e-4, (now - tF_prev_).seconds());
    dFdt = (Fe - Fe_prev_) / dts;
  }
  Fe_prev_ = Fe; tF_prev_ = now; has_prevF_ = true;

  const double Fenter = g_f_enter.load();
  const double Fexit  = g_f_exit.load();
  const double dFenter= g_df_enter.load();

  if (!in_contact_) {
    if (std::fabs(Fe) >= Fenter || std::fabs(dFdt) >= dFenter) {
      in_contact_ = true;
      // 初始化 TD：从当前实际力开始，向目标力平滑过渡
      g_td_v1 = Fe; g_td_v2 = 0.0; g_td_inited = true;
    }
  } else { // 已在接触
    if (std::fabs(Fe) <= Fexit && std::fabs(dFdt) <= 0.5 * dFenter) {
      in_contact_ = false; // 退出接触
      g_td_inited = false; // 下次进入接触再初始化
    }
  }

  // ======================
  // >>> TD/CONTACT：仅在接触阶段对 Fref 做 TD 平滑
  // ======================
  double Fref_eff = Fref; // 默认用原目标力
  if (in_contact_ && g_td_enable.load()) {
    const double r0 = g_td_r0.load();
    if (!g_td_inited) {                  // 双保险
      g_td_v1 = Fe; g_td_v2 = 0.0; g_td_inited = true;
    }
    td_update(Fref, r0, dt, g_td_v1, g_td_v2);
    Fref_eff = g_td_v1; // 用平滑后的设定力
  }

  // 线性映射为控制量 u（不使用 ESO/非线性反馈）
  // u 的物理意义：施加到“虚拟二阶系统”的驱动力
  const double eF = Fref_eff - Fe;
  const double u  = kappa * eF;

  // 阻抗模型：m_d * x¨ + b_d * x˙ + k_d * x = u
  const double md = std::max(1e-6, g_md.load());
  const double bd = g_bd.load();
  const double kd = g_kd.load();

  double xdd = (u - bd * xdot_ - kd * x_) / md;

  // 限制加速度
  const double Amax = std::max(0.0, g_max_acc.load());
  if (xdd >  Amax) xdd =  Amax;
  if (xdd < -Amax) xdd = -Amax;

  // 积分得到速度并限速
  double xdot_new = xdot_ + xdd * dt;
  const double Vmax = std::max(0.0, g_max_vel.load());
  if (xdot_new >  Vmax) xdot_new =  Vmax;
  if (xdot_new < -Vmax) xdot_new = -Vmax;

  xdot_ = xdot_new;
  x_   += xdot_ * dt;
  u_prev_ = u;

  // 发布 Twist（将 x˙ 映射到所选轴向）
  auto msg = std::make_unique<geometry_msgs::msg::TwistStamped>();
  msg->header.stamp = now;
  msg->header.frame_id = g_twist_frame;

  const int axis = g_axis.load();
  msg->twist.linear.x = (axis == 0) ? xdot_ : 0.0;
  msg->twist.linear.y = (axis == 1) ? xdot_ : 0.0;
  msg->twist.linear.z = (axis == 2) ? xdot_ : 0.0;
  msg->twist.angular.x = 0.0;
  msg->twist.angular.y = 0.0;
  msg->twist.angular.z = 0.0;

  twist_pub_->publish(std::move(msg));

  // 节流打印 —— 增加 TD/CONTACT 信息，其他不变
  RCLCPP_INFO_THROTTLE(
      LOGGER, *node_->get_clock(), 50,
      "[IMP] mode=%s Fe=%.3fN Fref_eff=%.3fN e=%.3f u=%.3f | x=%.4f xdot=%.4f (dt=%.4f) "
      "[md=%.2f bd=%.2f kd=%.2f | axis=%d Vmax=%.2f Amax=%.2f frame=%s | TD:%s r0=%.1f v2=%.3f dFdt=%.2f]",
      in_contact_ ? "CONTACT" : "FREE",
      Fe, Fref_eff, eF, u, x_, xdot_, dt,
      md, bd, kd, axis, Vmax, Amax, g_twist_frame.c_str(),
      g_td_enable.load() ? "ON" : "OFF",
      g_td_r0.load(), g_td_v2, dFdt);
}

} // namespace imp_node

int main(int argc, char** argv)
{
  using namespace imp_node;

  rclcpp::init(argc, argv);
  rclcpp::NodeOptions opts;
  opts.use_intra_process_comms(false);
  node_ = std::make_shared<rclcpp::Node>("plane_impedance", opts);

  // —— MoveIt Planning Scene Monitor（让 Servo 正常工作）——
  auto tf_buffer = std::make_shared<tf2_ros::Buffer>(node_->get_clock());
  auto psm = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
      node_, "robot_description", tf_buffer, "planning_scene_monitor");
  if (psm->getPlanningScene()) {
    psm->startStateMonitor("/joint_states");
    psm->setPlanningScenePublishingFrequency(25);
    psm->startPublishingPlanningScene(
        planning_scene_monitor::PlanningSceneMonitor::UPDATE_SCENE,
        "/moveit_servo/publish_planning_scene");
    psm->startSceneMonitor();
    psm->providePlanningSceneService();
  } else {
    RCLCPP_ERROR(LOGGER, "Planning scene 未正确配置");
    return EXIT_FAILURE;
  }

  // —— MoveIt Servo ——（必须创建并 start）
  auto servo_params = moveit_servo::ServoParameters::makeServoParameters(node_);
  auto servo = std::make_unique<moveit_servo::Servo>(node_, servo_params, psm);
  servo->start();

  // —— Pub/Sub ——
  twist_pub_ = node_->create_publisher<geometry_msgs::msg::TwistStamped>(
      "~/delta_twist_cmds", 10);

  force_sub_ = node_->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/sim_force", rclcpp::SensorDataQoS(), &forceCallback);

  // —— 参数声明 ——（都可 ros2 param set 动态修改）
  node_->declare_parameter<int>("axis", g_axis.load());
  node_->declare_parameter<double>("f_target", g_f_target.load());
  node_->declare_parameter<double>("kappa", g_kappa.load());

  node_->declare_parameter<double>("imp.md", g_md.load());
  node_->declare_parameter<double>("imp.bd", g_bd.load());
  node_->declare_parameter<double>("imp.kd", g_kd.load());

  node_->declare_parameter<double>("limit.max_vel", g_max_vel.load());
  node_->declare_parameter<double>("limit.max_acc", g_max_acc.load());

  node_->declare_parameter<std::string>("twist_frame", g_twist_frame);

  // >>> TD/CONTACT：新增参数声明
  node_->declare_parameter<double>("contact.f_enter", g_f_enter.load());
  node_->declare_parameter<double>("contact.f_exit",  g_f_exit.load());
  node_->declare_parameter<double>("contact.df_enter", g_df_enter.load());
  node_->declare_parameter<bool>("td.enable",   g_td_enable.load());
  node_->declare_parameter<double>("td.r0",     g_td_r0.load());

  // —— 同步初值 ——
  g_axis.store(node_->get_parameter("axis").as_int());
  g_f_target.store(node_->get_parameter("f_target").as_double());
  g_kappa.store(node_->get_parameter("kappa").as_double());

  g_md.store(node_->get_parameter("imp.md").as_double());
  g_bd.store(node_->get_parameter("imp.bd").as_double());
  g_kd.store(node_->get_parameter("imp.kd").as_double());

  g_max_vel.store(node_->get_parameter("limit.max_vel").as_double());
  g_max_acc.store(node_->get_parameter("limit.max_acc").as_double());

  g_twist_frame = node_->get_parameter("twist_frame").as_string();

  // >>> TD/CONTACT：同步新增参数
  g_f_enter.store(node_->get_parameter("contact.f_enter").as_double());
  g_f_exit.store( node_->get_parameter("contact.f_exit").as_double());
  g_df_enter.store(node_->get_parameter("contact.df_enter").as_double());
  g_td_enable.store(node_->get_parameter("td.enable").as_bool());
  g_td_r0.store(    node_->get_parameter("td.r0").as_double());

  // —— 动态参数回调 ——
  param_cb_handle = node_->add_on_set_parameters_callback(
    [](const std::vector<rclcpp::Parameter>& params)
    {
      rcl_interfaces::msg::SetParametersResult res; res.successful = true;
      for (const auto& p : params) {
        const auto &name = p.get_name();
        if      (name == "axis")             { imp_node::g_axis.store(p.as_int()); }
        else if (name == "f_target")         { imp_node::g_f_target.store(p.as_double()); }
        else if (name == "kappa")            { imp_node::g_kappa.store(p.as_double()); }
        else if (name == "imp.md")           { imp_node::g_md.store(p.as_double()); }
        else if (name == "imp.bd")           { imp_node::g_bd.store(p.as_double()); }
        else if (name == "imp.kd")           { imp_node::g_kd.store(p.as_double()); }
        else if (name == "limit.max_vel")    { imp_node::g_max_vel.store(p.as_double()); }
        else if (name == "limit.max_acc")    { imp_node::g_max_acc.store(p.as_double()); }
        else if (name == "twist_frame")      { imp_node::g_twist_frame = p.as_string(); }

        // >>> TD/CONTACT：动态更新
        else if (name == "contact.f_enter")  { imp_node::g_f_enter.store(p.as_double()); }
        else if (name == "contact.f_exit")   { imp_node::g_f_exit.store(p.as_double()); }
        else if (name == "contact.df_enter") { imp_node::g_df_enter.store(p.as_double()); }
        else if (name == "td.enable")        { imp_node::g_td_enable.store(p.as_bool()); }
        else if (name == "td.r0")            { imp_node::g_td_r0.store(p.as_double()); }
      }
      return res;
    }
  );

  // —— 控制循环：10ms（100Hz） ——
  auto timer = node_->create_wall_timer(10ms, &controlTimer);

  auto exec = std::make_unique<rclcpp::executors::MultiThreadedExecutor>();
  exec->add_node(node_);
  exec->spin();

  rclcpp::shutdown();
  return 0;
}
